package dan.frontend
import  chisel3._
import chisel3.util._
import  dan.utility._
import chisel3.util.random.LFSR
import dan.utility.FrontendUtility.fetchAlign

// 2位计数器的状态机
class BiModelState extends Bundle{
    val state = UInt(2.W)
    def isTaken = state(0)
    def init(taken: Bool) = {
        state := Mux(taken, 1.U, 0.U)
    }
    def update(taken: Bool) = {
        val next = Cat(
            taken ^ state(0),
            (!taken & !state(1) & state(0)) | (taken & (state(1) | state(0)))
        )
        this.state := next
    }
}
// 跳转地址
class MicroBTBEntry extends Bundle with HasMicroBTBParam{
    val target = UInt(targetBits.W)
}
// tag，2位计数器等辅助信息
class MicroBTBInfo extends Bundle with HasMicroBTBParam{
    val tag = UInt(tagBits.W)
    val isBr = Bool()
    val state = new BiModelState()
}
// 预测信息
class MicroBTBMeta extends Bundle with HasMicroBTBParam{
    val hit = Bool()
    val wayIdx = UInt(log2Ceil(wayNum).W)
}

class MicroBTBIO extends Bundle with HasMicroBTBParam{
    val s1Valid = Input(Bool())
    val s1PC = Input(UInt(vaBits.W))
    val s1HashedPC = Input(UInt(vaBits.W))
    val s1Target = Output(Vec(fetchWidth, Valid(UInt(targetBits.W))))
    val s1Taken = Output(Vec(fetchWidth, Bool()))
    val s1IsBr = Output(Vec(fetchWidth, Bool()))
    val s1IsJal = Output(Vec(fetchWidth, Bool()))
    val s3Meta = Output(Vec(fetchWidth, new MicroBTBMeta()))
    val s1Update = Input(Valid(new PredictionUpdate()))
}

class MicroBTB extends  Module with HasMicroBTBParam{
    val io = IO(new MicroBTBIO())
    val valid = Reg(Vec(wayNum*fetchWidth, Bool()))
    val info = Reg(Vec(wayNum, Vec(fetchWidth, new MicroBTBInfo())))
    val targets = Reg(Vec(wayNum, Vec(fetchWidth, new MicroBTBEntry())))

    val s1Tag = getTag(io.s1HashedPC)
    val s1HitVec = VecInit((0 until fetchWidth) map{
        i => VecInit((0 until wayNum) map {
            // tag相等且表项有效
            j => IsEqual(info(j)(i).tag, s1Tag) && valid(Cat(j.U(log2Ceil(wayNum).W), i.U(log2Ceil(fetchWidth).W)))
        })
    })
    // 当前取值包中每条指令的命中信息
    val s1GetHit = s1HitVec.map(_.asUInt.orR)
    val s1HitWay = s1HitVec.map(vec => OHToUInt(vec))
    // 读出预测结果
    for (i <- 0 until fetchWidth){
        // 命中与否
        val giveResp = io.s1Valid && s1GetHit(i)
        io.s1Target(i).valid := giveResp
        // 注意shape是 wayNum x fetchWidth
        io.s1Target(i).bits := (targets(s1HitWay(i))(i).target).asUInt

        val entryInfo = info(s1HitWay(i))(i)
        io.s1IsBr(i) := giveResp && entryInfo.isBr
        io.s1IsJal(i) := giveResp && !entryInfo.isBr
        // 预测跳转或者是jal指令
        io.s1Taken(i) := entryInfo.state.isTaken || !entryInfo.isBr
    }

    // 如果没有命中，准备随机替换
    val replaceEnable = io.s1Valid && !s1GetHit.reduce(_||_)
    // 产生一个随机数
    val replaceWay = LFSR(wayNum, replaceEnable)(log2Ceil(wayNum)-1, 0)
    val s1Meta = Wire(Vec(fetchWidth, new MicroBTBMeta()))
    for(i <- 0 until fetchWidth) {
        s1Meta(i).hit := s1GetHit(i)
        // 如果没有命中，则目标way为随机替换的
        s1Meta(i).wayIdx := Mux(s1GetHit(i), s1HitWay(i), replaceWay)
    }
    io.s3Meta := RegNext(RegNext(s1Meta))

    val s1Update = io.s1Update
    // 更新的分支指令索引
    val s1UpdateCFIIdx = s1Update.bits.idx.bits
    // 每条指令的meta（命中）信息
    val s1UpdateMeta = VecInit(s1Update.bits.meta.map(_.microBTBMeta))
    val s1UpdateWayIdx = VecInit(s1UpdateMeta.map(_.wayIdx))
    // 最终需要更新的way索引
    val s1UpdateFinalWay = s1UpdateWayIdx(s1UpdateCFIIdx)
    // 允许更新
    val wrEnable = s1Update.valid && s1Update.bits.taken && s1Update.bits.idx.valid
    // 更新目标地址
    // 一个way存储同一取值组的4条指令
    when(wrEnable){
        targets(s1UpdateFinalWay)(s1UpdateCFIIdx).target := getTarget(s1Update.bits.target)
    }
    val taken = WireInit(VecInit(Seq.fill(fetchWidth)(false.B)))
    // 更新2位计数器状态和info，注意：所有可能的分支指令都需要更新
    for(i <- 0 until fetchWidth){
        val brUpdatable = s1Update.valid && ( s1Update.bits.brMask(i) || 
            // 这个条件冗余 ？
            (s1Update.bits.idx.valid && s1UpdateCFIIdx===i.U && s1Update.bits.taken)
        )
        when(brUpdatable){
            // isJal冗余 ？
            taken(i) := s1Update.bits.idx.valid && s1UpdateCFIIdx===i.U && (s1Update.bits.taken || s1Update.bits.isJal)
            // 同一取值组使用相同的tag
            valid(Cat(s1UpdateFinalWay, i.U(log2Ceil(fetchWidth).W))) := true.B
            info(s1UpdateFinalWay)(i).tag := getTag(mixHILO(s1Update.bits.pc))
            info(s1UpdateFinalWay)(i).isBr := s1Update.bits.brMask(i)

            // 更新2位计数器
            when(s1UpdateMeta(i).hit){
                // 如果以前命中，则存在一个计数器，更新状态
                info(s1UpdateFinalWay)(i).state.update(taken(i))
            }.otherwise{
                // 如果没有命中，新分配一项
                info(s1UpdateFinalWay)(i).state.init(taken(i))
            }
        }
    }
}

